import torch.nn as nn
import torch


class ConvLSTMCell(nn.Module):

    def __init__(self, input_dim, hidden_dim, kernel_size, bias=True):
        """
        初始化卷积 LSTM 单元。

        参数:
        ----------
        input_dim: int
            输入张量的通道数。
        hidden_dim: int
            隐藏状态的通道数。
        kernel_size: (int, int)
            卷积核的大小。
        bias: bool
            是否添加偏置项。
        """

        super(ConvLSTMCell, self).__init__()

        self.input_dim = input_dim
        self.hidden_dim = hidden_dim

        self.kernel_size = kernel_size
        # 计算填充大小以保持输入和输出尺寸一致
        self.padding = kernel_size[0] // 2, kernel_size[1] // 2
        self.bias = bias

        # 定义卷积层，输入是输入维度加上隐藏维度，输出是4倍的隐藏维度（对应i, f, o, g）
        self.conv = nn.Conv2d(in_channels=self.input_dim + self.hidden_dim,
                              out_channels=4 * self.hidden_dim,
                              kernel_size=self.kernel_size,
                              padding=self.padding,
                              bias=self.bias)

    def forward(self, input_tensor, cur_state):
        h_cur, c_cur = cur_state

        # 沿着通道轴进行拼接
        combined = torch.cat([input_tensor, h_cur], dim=1)

        combined_conv = self.conv(combined)
        # 将输出分割成四个部分，分别对应输入门、遗忘门、输出门和候选单元状态
        cc_i, cc_f, cc_o, cc_g = torch.split(combined_conv, self.hidden_dim, dim=1)
        i = torch.sigmoid(cc_i)
        f = torch.sigmoid(cc_f)
        o = torch.sigmoid(cc_o)
        g = torch.tanh(cc_g)

        # 更新单元状态
        c_next = f * c_cur + i * g
        # 更新隐藏状态
        h_next = o * torch.tanh(c_next)

        return h_next, c_next

    def init_hidden(self, batch_size, image_size):
        height, width = image_size
        # 初始化隐藏状态和单元状态为零
        return (torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device),
                torch.zeros(batch_size, self.hidden_dim, height, width, device=self.conv.weight.device))

